from sympy import symbols, simplify, expand, Eq, solve, root, sqrt, powsimp
from sympy.abc import a, b, c, d, e, f, x, t
import sympy as sp

# 安全的三次根实现
def safe_cbrt(expr):
    return root(expr, 3)

# 修正后的五次方程求解
def quintic_roots(a_val, b_val, c_val, d_val, e_val, f_val):
    # 转换为符号计算
    a, b, c, d, e, f = [sp.sympify(v) for v in [a_val, b_val, c_val, d_val, e_val, f_val]]
    
    # 1. 消去四次项
    alpha = (2*b**2)/(5*a**2) - c/a
    r = (2*b**3)/(25*a**3) - d/a
    Pi = (b**4)/(125*a**4) - e/a
    M = f/a - (b**5)/(3125*a**5)
    
    # 2. 计算动态参数m
    m = M - (b*Pi)/(5*a)
    
    # 3. 构造判别式Δ
    d_new = -alpha
    p = -(r - (3*b*alpha)/(5*a))
    q = -((3*b**2*alpha)/(25*a**2) - (2*b*r)/(5*a) - Pi)
    
    numerator = (m**2 + q*p**2 - d_new*m*p)/(2*m)
    term = (-(p + q)/3 - (d_new**2)/9)**3
    Delta = simplify(numerator**2 + term)
    
    # 4. 计算关键量A（使用安全的三次根）
    sqrt_Delta = sqrt(Delta)
    A = simplify(safe_cbrt(numerator + sqrt_Delta) + safe_cbrt(numerator - sqrt_Delta) - d_new/3)
    
    # 5. 因式分解
    D = -A
    # 显式求解方程组
    E = symbols('E')
    B = A**2 - d_new - E
    C = m/E
    eq = Eq(C + A*E + B*D, p)
    E_sol = solve(eq, E)
    
    # 6. 求解低次方程
    solutions = []
    for E_val in E_sol:
        B_val = A**2 - d_new - E_val
        C_val = m/E_val
        
        # 三次部分
        cubic = x**3 + A*x**2 + B_val*x + C_val
        cubic_roots = sp.solve(cubic, x)
        
        # 二次部分
        quadratic = x**2 + D*x + E_val
        quad_roots = sp.solve(quadratic, x)
        
        # 合并解
        solutions.extend(cubic_roots + quad_roots)
    
    # 7. 回代
    final_roots = [simplify(root - b/(5*a)) for root in solutions]
    return final_roots

# 数值验证案例
def numerical_verification():
    test_cases = [
        (1, 0, -5, 0, 5, -2),  # x^5 -5x^3 +5x -2 = 0
        (1, 0, 0, 0, -2, 1),    # x^5 -2x +1 = 0
        (1, -5, 10, -10, 5, -1) # (x-1)^5 = 0
    ]
    
    for coeffs in test_cases:
        print(f"\n验证方程: {coeffs[0]}x^5 + {coeffs[1]}x^4 + ... + {coeffs[5]} = 0")
        roots = quintic_roots(*coeffs)
        
        for i, r in enumerate(roots):
            # 构建多项式
            P = sum(c*x**i for i, c in enumerate(reversed(coeffs)))
            # 代入验证
            val = simplify(P.subs(x, r))
            print(f"根x_{i+1} = {r}")
            print(f"残差: {val}")
            assert val == 0, "验证失败！"

print("开始数值验证...")
numerical_verification()
print("基本测试案例验证通过！")

# 符号验证（简化版）
def symbolic_proof():
    P = a*x**5 + b*x**4 + c*x**3 + d*x**2 + e*x + f
    print("\n开始符号验证...")
    
    try:
        roots = quintic_roots(a, b, c, d, e, f)
        for i, r in enumerate(roots[:2]):  # 只验证前两个根避免过载
            print(f"\n验证符号根 {i+1}:")
            print("根表达式:", r)
            res = simplify(P.subs(x, r))
            print("代入结果:", res)
            print("展开尝试:", expand(res))
    except Exception as e:
        print(f"符号验证遇到限制: {str(e)}")
        print("建议使用数值验证或特定系数验证")

symbolic_proof()